#include "ns3/core-module.h"
#include "ns3/function-mobility.h"
#include "ns3/csma-module.h"
#include "ns3/internet-module.h"
#include "ns3/network-module.h"
#include "ns3/openflow-module.h"
#include "ns3/robo-module.h"

using namespace ns3;
using namespace std;

struct UploadInfo
{
  uint8_t cmd;
  uint8_t uid;
  FVector loc;
  int num17;
  int heat17;
  int num42;
  int heat42;
};

struct LocInfo
{
  uint8_t cmd;
  uint8_t uid;
  FVector loc;
  int32_t ts;
};

struct Moveto
{
  uint8_t cmd;
  FVector target;
};

struct Shoot
{
  uint8_t cmd;
  AmmoType type;
  FVector loc;
};

class RoboFoolishAppClient : public Application
{
public:
  virtual void StartApplication (void);
  virtual void StopApplication (void);
  void Update (float delta);
  void IndicateLocCb (uint8_t uid, FVector loc);
  void HandleCollision (Ptr<RoboActor> oth);
  void HandleRecv (Ptr<Socket> socket);

  Ptr<Socket> m_socket;
  Address m_peer;

  RoboType m_roboType;
  Ptr<RoboBase> m_robo;
  Ptr<VariableSpeedMobility> m_mobv;
  Ptr<FunctionMobility> m_mobf;
  bool m_bStopped{false};
  Time nxtSend;
  // FVector m_target;
  // std::map<uint8_t, std::pair<FVector, Time>> m_knownLocation;

public:
  static TypeId GetTypeId (void);
};

TypeId
RoboFoolishAppClient::GetTypeId (void)
{
  static TypeId tid = TypeId ("ns3::RoboFoolishAppClient")
                          .SetParent<Application> ()
                          .SetGroupName ("Robo")
                          .AddConstructor<RoboFoolishAppClient> ();
  return tid;
}

void
RoboFoolishAppClient::StartApplication (void)
{
  if (!m_socket)
    {
      m_socket = Socket::CreateSocket (GetNode (), UdpSocketFactory::GetTypeId ());
      if (Inet6SocketAddress::IsMatchingType (m_peer))
        {
          if (m_socket->Bind6 () == -1)
            {
              NS_FATAL_ERROR ("Failed to bind socket");
            }
        }
      else if (InetSocketAddress::IsMatchingType (m_peer) ||
               PacketSocketAddress::IsMatchingType (m_peer))
        {
          if (m_socket->Bind () == -1)
            {
              NS_FATAL_ERROR ("Failed to bind socket");
            }
        }
      m_socket->Connect (m_peer);
      m_socket->SetAllowBroadcast (true);
      m_socket->SetRecvCallback (MakeCallback (&RoboFoolishAppClient::HandleRecv, this));
    }

  m_robo = m_node->GetObject<RoboBase> ();
  NS_ASSERT_MSG (m_robo, "No robo");
  m_mobv = DynamicCast<VariableSpeedMobility> (m_robo->GetMobility ());
  m_mobf = DynamicCast<FunctionMobility> (m_robo->GetMobility ());
  if (m_mobf)
    m_mobf->SetSpeed (1.f);
  NS_ASSERT_MSG (m_robo->GetMobility (), "No mobility");
  m_roboType = m_robo->GetRoboType ();
  m_robo->AddUpdateCallback (MakeCallback (&RoboFoolishAppClient::Update, this));
  m_robo->AddIndicateLocationCallback (MakeCallback (&RoboFoolishAppClient::IndicateLocCb, this));
  m_robo->AddCollisionCallback (MakeCallback (&RoboFoolishAppClient::HandleCollision, this));
}
void
RoboFoolishAppClient::StopApplication (void)
{
  m_bStopped = true;
}

void
RoboFoolishAppClient::IndicateLocCb (uint8_t uid, FVector loc)
{
  LocInfo info;
  info.cmd = 2;
  info.uid = uid;
  info.loc = loc;
  info.ts = Simulator::Now ().GetMicroSeconds ();
  uint8_t *buffer = (uint8_t *) ((void *) (&info));
  Ptr<Packet> pack = Create<Packet> (buffer, sizeof (LocInfo));
  m_socket->Send (pack);
}

void
RoboFoolishAppClient::HandleCollision (Ptr<RoboActor> oth)
{
  Ptr<RoboSupply> sup = DynamicCast<RoboSupply> (oth);
  if (sup)
    {
      sup->AskForAmmo (m_robo, 10, AmmoType17);
    }
}

void
RoboFoolishAppClient::HandleRecv (Ptr<Socket> socket)
{
  const int bufferSize = 1024;
  static uint8_t buffer[bufferSize];
  Ptr<Packet> packet;
  Address from;
  while ((packet = socket->RecvFrom (from)))
    {
      if (packet->GetSize () == 0)
        { //EOF
          break;
        }
      packet->CopyData (buffer, bufferSize);
      switch (*buffer)
        {
        case 3:
          {
            Moveto *move = (Moveto *) ((void *) (buffer));
            auto diff = move->target - m_robo->GetGlobalLocation ();
            if (diff.GetLength () > .5f)
              {
                auto speed = diff.Norm () * gRoboMaxSpeed[m_robo->GetRoboType ()];
                m_mobv->SetSpeed (speed);
              }
            else
              {
                m_mobv->SetSpeed (FVector (0.f, 0.f));
              }
            break;
          }
        case 4:
          {
            Shoot *shoot = (Shoot *) ((void *) (buffer));
            auto diff = shoot->loc - m_robo->GetGlobalLocation ();
            auto speed = diff.Norm () * 7;
            if (shoot->type == AmmoType42)
              m_robo->ShootAmmo (speed, AmmoType42);
            m_robo->ShootAmmo (speed, AmmoType17);
            break;
          }

        default:
          break;
        }
    }
}

void
RoboFoolishAppClient::Update (float delta)
{
  if (m_bStopped || Simulator::Now () < nxtSend)
    return;
  nxtSend = Simulator::Now () + Seconds (0.1);
  UploadInfo info;
  info.cmd = 1;
  info.uid = m_robo->GetUid ();
  info.loc = m_robo->GetGlobalLocation ();
  info.num17 = m_robo->GetAmmoNum (AmmoType17);
  info.heat17 = m_robo->GetHeat (AmmoType17);
  info.num42 = m_robo->GetAmmoNum (AmmoType42);
  info.heat42 = m_robo->GetHeat (AmmoType42);

  uint8_t *buffer = (uint8_t *) ((void *) (&info));
  Ptr<Packet> pack = Create<Packet> (buffer, sizeof (UploadInfo));
  m_socket->Send (pack);
}

class RoboFoolishAppServer : public Application
{
public:
  FVector m_target;
  uint8_t m_team;
  Address m_local;
  Ptr<Socket> m_socket;
  Ptr<Socket> m_sendSocket;
  std::map<uint8_t, std::pair<FVector, Time>> m_knownLocation;
  static TypeId GetTypeId (void);
  virtual void StartApplication (void);
  virtual void StopApplication (void);
  void HandleRead (Ptr<Socket> socket);
};

TypeId
RoboFoolishAppServer::GetTypeId (void)
{
  static TypeId tid = TypeId ("ns3::RoboFoolishAppServer")
                          .SetParent<Application> ()
                          .SetGroupName ("Robo")
                          .AddConstructor<RoboFoolishAppServer> ();
  return tid;
}

void
RoboFoolishAppServer::HandleRead (Ptr<Socket> socket)
{
  const int bufferSize = 1024;
  static uint8_t buffer[bufferSize];
  Ptr<Packet> packet;
  Address from;
  while ((packet = socket->RecvFrom (from)))
    {
      if (packet->GetSize () == 0)
        { //EOF
          break;
        }
      packet->CopyData (buffer, bufferSize);
      // InetSocketAddress::ConvertFrom (from).GetIpv4 ();
      switch (*buffer)
        {
        case 1:
          {
            UploadInfo *info = (UploadInfo *) ((void *) (buffer));
            auto roboStatus = RoboJudge::Get ()->QueryRobo (info->uid);
            switch (roboStatus.type)
              {
              case RoboTypeHero:
              case RoboTypeInfantry:
              case RoboTypeEngineer:
              case RoboTypeUVA:
                {
                  if (info->loc.GetDistance2 (m_target) > 4.f)
                    {
                      Moveto move;
                      move.cmd = 3;
                      move.target = m_target;
                      Ptr<Packet> pack =
                          Create<Packet> ((uint8_t *) ((void *) (&move)), sizeof (Moveto));
                      socket->SendTo (pack, 0, from);
                    }
                  break;
                }

              default:
                break;
              }
            break;
          }
        case 2:
          {
            LocInfo *info = (LocInfo *) ((void *) (buffer));
            m_knownLocation[info->uid] = std::make_pair (info->loc, MicroSeconds (info->ts));
            auto roboStatus = RoboJudge::Get ()->QueryRobo (info->uid);
            if (roboStatus.team != m_team)
              {
                Shoot shoot;
                shoot.cmd = 4;
                shoot.type = AmmoType42;
                shoot.loc = info->loc;
                Ptr<Packet> pack = Create<Packet> ((uint8_t *) ((void *) (&shoot)), sizeof (Shoot));
                socket->SendTo (pack, 0, from);
              }
            break;
          }

        default:
          std::cerr << "Unknow code " << (int) *buffer << std::endl;
        }
    }
}

void
RoboFoolishAppServer::StartApplication (void)
{
  if (!m_socket)
    {
      m_socket = Socket::CreateSocket (GetNode (), UdpSocketFactory::GetTypeId ());
      if (m_socket->Bind (m_local) == -1)
        {
          NS_FATAL_ERROR ("Failed to bind socket");
        }
      m_socket->Listen ();
    }
  m_socket->SetRecvCallback (MakeCallback (&RoboFoolishAppServer::HandleRead, this));
}
void
RoboFoolishAppServer::StopApplication (void)
{
  if (m_socket)
    {
      m_socket->Close ();
    }
}

int
main (int argc, char **argv)
{
#ifdef NS3_OPENFLOW

  // LogComponentEnable ("RoboBase", LOG_LEVEL_DEBUG);
  // LogComponentEnable ("OpenFlowInterface", LOG_LEVEL_INFO);
  // LogComponentEnable ("OpenFlowSwitchNetDevice", LOG_LEVEL_INFO);
  // Config::SetDefault ("ns3::RoboJudge::EnablePlot", BooleanValue (false));

  RoboJudge::Get ()->Init ();
  RoboTeamConfig cfg;
  cfg.push_back ({RoboTypeHero, "Hero", FVector (2, 1), FAngle (0)});
  cfg.push_back ({RoboTypeInfantry, "Infantry1", FVector (1, 10), FAngle (0)});
  cfg.push_back ({RoboTypeInfantry, "Infantry2", FVector (10, 1), FAngle (0)});
  cfg.push_back ({RoboTypeSentinel, "Sentinel", FVector (0, 0), FAngle (0)});
  cfg.push_back ({RoboTypeEngineer, "Engineer", FVector (1, 2), FAngle (0)});
  cfg.push_back ({RoboTypeUVA, "UVA", FVector (5, 0.5), FAngle (0)});
  cfg.push_back ({RoboTypeTower, "Tower", FVector (1., 1.), FAngle (0)});
  RoboTeamHelper team1;
  team1.Init (1, FVector (0, 0), FAngle (0), cfg);
  RoboTeamHelper team2;
  team2.Init (2, FVector (28, 15), FAngle (PI), cfg);

  NodeContainer nodes1 = team1.GetNodes ();
  NodeContainer nodes2 = team2.GetNodes ();
  NodeContainer nodes = NodeContainer (nodes1, nodes2);

  /////////////////////////////////////////////////Network////////////////////////////////////////////////////////////////////
  NodeContainer ofSwitch;
  ofSwitch.Create (1);

  CsmaHelper csma;
  csma.SetChannelAttribute ("DataRate", DataRateValue (5000000));
  csma.SetChannelAttribute ("Delay", TimeValue (MilliSeconds (2)));

  NetDeviceContainer terminalDevices;
  NetDeviceContainer switchDevices;
  for (uint32_t i = 0; i < nodes.GetN (); i++)
    {
      NetDeviceContainer link = csma.Install (NodeContainer (nodes.Get (i), ofSwitch));
      terminalDevices.Add (link.Get (0));
      switchDevices.Add (link.Get (1));
    }
  Ptr<Node> switchNode = ofSwitch.Get (0);
  OpenFlowSwitchHelper swtch;

  Ptr<ns3::ofi::LearningController> controller = CreateObject<ns3::ofi::LearningController> ();
  swtch.Install (switchNode, switchDevices, controller);

  InternetStackHelper internet;
  internet.Install (nodes);

  Ipv4AddressHelper ipv4;
  ipv4.SetBase ("10.1.1.0", "255.255.255.0");
  ipv4.Assign (terminalDevices);

  for (uint32_t i = 0; i < nodes1.GetN (); ++i)
    {
      Ptr<RoboBase> robo = nodes1.Get (i)->GetObject<RoboBase> ();
      if (robo->GetRoboType () == RoboTypeTower)
        continue;
      // std::cerr << "Install client app on " << robo->GetName () << std::endl;
      Ptr<RoboFoolishAppClient> app = CreateObject<RoboFoolishAppClient> ();
      app->m_peer = Address (InetSocketAddress (Ipv4Address ("10.1.1.7"), 9));
      nodes1.Get (i)->AddApplication (app);
    }
  Ptr<RoboFoolishAppServer> app1 = CreateObject<RoboFoolishAppServer> ();
  app1->m_local = Address (InetSocketAddress (Ipv4Address ("10.1.1.7"), 9));
  app1->m_target = FVector (26, 13);
  app1->m_team = 1;
  team1.GetNodeOf ("Tower")->AddApplication (app1);

  for (uint32_t i = 0; i < nodes2.GetN (); ++i)
    {
      Ptr<RoboBase> robo = nodes2.Get (i)->GetObject<RoboBase> ();
      if (robo->GetRoboType () == RoboTypeTower)
        continue;
      // std::cerr << "Install client app on " << robo->GetName () << std::endl;
      Ptr<RoboFoolishAppClient> app = CreateObject<RoboFoolishAppClient> ();
      app->m_peer = Address (InetSocketAddress (Ipv4Address ("10.1.1.14"), 9));
      nodes2.Get (i)->AddApplication (app);
    }
  Ptr<RoboFoolishAppServer> app2 = CreateObject<RoboFoolishAppServer> ();
  app2->m_local = Address (InetSocketAddress (Ipv4Address ("10.1.1.14"), 9));
  app2->m_target = FVector (2, 2);
  app2->m_team = 2;
  team2.GetNodeOf ("Tower")->AddApplication (app2);

  NS_LOG_UNCOND("Simulation start...");
  Simulator::Stop (Seconds (60));
  Simulator::Run ();
  Simulator::Destroy ();
  NS_LOG_UNCOND("End");

#else
  NS_ABORT_MSG ("NS-3 OpenFlow is not enabled. Cannot run simulation.");
#endif
  return 0;
}